import pywt
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams
from pywt import wavelist

# 设置中文字体
rcParams['font.sans-serif'] = ['SimHei']  # Windows系统黑体
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 1. 读取期货时间序列（示例数据格式）
df = pd.read_csv('data\SF0SH.csv',index_col='date_time')
df.index = pd.to_datetime(df.index)  # 确保时间索引为datetime类型
prices = df['close'].values
timestamps = df.index  # 获取时间序列

# 2. 小波分解参数设置
wavelet = 'sym6'       # 选用Daubechies4小波基[[19]]
level = 3             # 分解层数建议为3层[[21]]
mode = 'symmetric'    # 边界对称延拓处理(symmetric padding)
                     # 可选模式: ['zero', 'constant', 'symmetric', 'periodic', 'smooth', 'reflect']

# 3. 执行小波分解
coeffs = pywt.wavedec(prices, wavelet, level=level, mode=mode)
cA = coeffs[0]        # 低频逼近信号[[1]]
cD_list = coeffs[1:]  # 高频细节信号列表（从高层到低层）

# 4. 重构各层信号
def reconstruct_coeff(coeffs, level):
    """重构指定层的信号"""
    coeffs_temp = [np.zeros_like(c) for c in coeffs]
    coeffs_temp[level] = coeffs[level]
    return pywt.waverec(coeffs_temp, wavelet, mode=mode)

# 添加信号到可视化列表（原始序列在最上方）
recon_signals = [prices]  # 原始信号
recon_signals.append(cA)   # 低频信号
for i in range(level):
    # 重构第i+1层的细节信号（注意层级索引关系）
    recon_signals.append(reconstruct_coeff(coeffs, i+1))

# 5. 绘制对比图
plt.figure(figsize=(15, 18))  # 增加画布高度
titles = ['原始序列'] + ['低频信号'] + [f'第{level-i}层细节信号' for i in range(level)]

for i, sig in enumerate(recon_signals):
    plt.subplot(len(recon_signals), 1, i+1)  # 单列布局
    plt.plot(timestamps[:len(sig)], sig, 'b', linewidth=1)
    plt.title(titles[i], fontsize=10)
    ax = plt.gca()
    ax.xaxis.set_major_formatter(plt.matplotlib.dates.DateFormatter('%m-%d'))
    ax.xaxis.set_major_locator(plt.matplotlib.dates.DayLocator(interval=5))
    plt.xticks(rotation=30)
    plt.subplots_adjust(hspace=0.5, bottom=0.2)  # 增加底部边距

plt.tight_layout()
plt.show()
